package bloom_filter

import (
	"math"

	"github.com/spaolacci/murmur3"
)

const (
	DefaultFalsePositive = 1e-20
)

// BloomFilter represents a probabilistic data structure for membership testing
type BloomFilter struct {
	m        int    // number of bits
	k        int    // number of hash functions
	bitArray []byte // the bit array
	n        int    // number of elements
}

// NewBloomFilter creates a new BloomFilter with specified m and k
func NewBloomFilter(m, k int) *BloomFilter {
	numBytes := m / 8
	if m%8 != 0 {
		numBytes++
	}
	return &BloomFilter{
		m:        m,
		k:        k,
		bitArray: make([]byte, numBytes),
		n:        0,
	}
}

func (bf *BloomFilter) Empty() bool {
	return bf.n == 0
}

// getNthBit gets the value of the nth bit in the bit array
func (bf *BloomFilter) getNthBit(n int) byte {
	byteIndex := n / 8
	bitIndex := n % 8
	return (bf.bitArray[byteIndex] >> bitIndex) & 1
}

// setNthBit sets the nth bit in the bit array to 1
func (bf *BloomFilter) setNthBit(n int) {
	byteIndex := n / 8
	bitIndex := n % 8
	bf.bitArray[byteIndex] |= 1 << bitIndex
}

// nthHash calculates the nth hash value using double hashing
func (bf *BloomFilter) nthHash(n int, h1, h2 uint64) int {
	hash := int((h1 + uint64(n)*h2) % uint64(bf.m))
	if hash < 0 {
		hash += bf.m
	}
	return hash
}

// Add adds a key to the Bloom filter
func (bf *BloomFilter) Add(key []byte) {
	h1, h2 := murmur3.Sum128(key)
	for i := 0; i < bf.k; i++ {
		hash := bf.nthHash(i, h1, h2)
		bf.setNthBit(hash)
	}
	bf.n++
}

// MightContain checks if a key might be contained in the Bloom filter
func (bf *BloomFilter) MightContain(key []byte) bool {
	h1, h2 := murmur3.Sum128(key)
	for i := 0; i < bf.k; i++ {
		hash := bf.nthHash(i, h1, h2)
		if bf.getNthBit(hash) == 0 {
			return false
		}
	}
	return true
}

// FalsePositiveProbability calculates the current false positive probability
func (bf *BloomFilter) FalsePositiveProbability() float64 {
	return math.Pow(1-math.Exp(-1.0*float64(bf.k*bf.n)/float64(bf.m)), float64(bf.k))
}

// EnsureFalsePositiveProbabilityUnder calculates optimal m and k for given prob and n
func EnsureFalsePositiveProbabilityUnder(prob float64, n int64) (int, int) {
	// Calculate the optimal size of the bit array m
	m := int(math.Ceil(-1 * float64(n) * math.Log(prob) / (math.Log(2) * math.Log(2))))

	// Calculate the optimal number of hash functions k
	k := int(math.Ceil(float64(m) / float64(n) * math.Log(2)))

	return m, k
}

// Merge combines another Bloom filter into this one (A ∪ B)
func (bf *BloomFilter) Merge(another *BloomFilter) {
	for i := 0; i < len(bf.bitArray); i++ {
		bf.bitArray[i] |= another.bitArray[i]
	}
}

// Clone creates a deep copy of the Bloom filter
func (bf *BloomFilter) Clone() *BloomFilter {
	newBf := NewBloomFilter(bf.m, bf.k)
	newBf.n = bf.n
	copy(newBf.bitArray, bf.bitArray)
	return newBf
}
